#!/usr/bin/env python3
"""
重建FAISS索引工具
当索引损坏或不一致时，从数据库重建索引
"""
import sys
import os

# 添加项目根目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database.connection import get_sync_session
from repositories.orm_face_repository import OrmFaceRepository
from utils.faiss_index import faiss_manager
from utils.logger import app_logger
import numpy as np


def rebuild_index():
    """从数据库重建FAISS索引"""
    print("=" * 80)
    print("FAISS索引重建工具")
    print("=" * 80)
    
    try:
        # 1. 检查当前索引状态
        print(f"\n📊 当前索引状态:")
        print(f"   索引向量数量: {faiss_manager.index.ntotal if faiss_manager.index else 0}")
        print(f"   映射记录数量: {len(faiss_manager.face_id_map)}")
        
        is_valid = faiss_manager.validate_index()
        if is_valid:
            print(f"   ✅ 索引状态正常")
            response = input("\n索引正常，是否仍要重建？(y/N): ")
            if response.lower() != 'y':
                print("取消重建")
                return
        else:
            print(f"   ❌ 索引不一致，需要重建！")
        
        # 2. 从数据库加载所有人脸数据
        print(f"\n📥 从数据库加载人脸数据...")
        with get_sync_session() as session:
            face_repo = OrmFaceRepository(session)
            all_faces = face_repo.get_all_faces(include_deleted=False)
            
            if not all_faces:
                print("⚠️  数据库中没有人脸数据")
                faiss_manager._create_new_index()
                faiss_manager.save()
                print("✅ 创建空索引完成")
                return
            
            print(f"   找到 {len(all_faces)} 条人脸记录")
        
        # 3. 准备数据
        print(f"\n🔧 准备重建数据...")
        face_data = []
        for face in all_faces:
            try:
                # 解析特征向量
                feature_bytes = face.feature_vector
                feature_array = np.frombuffer(feature_bytes, dtype=np.float32)
                
                if feature_array.shape[0] != 512:
                    print(f"   ⚠️  跳过无效特征向量: {face.face_id} (维度: {feature_array.shape[0]})")
                    continue
                
                face_data.append((face.face_id, feature_array))
            except Exception as e:
                print(f"   ⚠️  跳过错误记录: {face.face_id}, 错误: {e}")
                continue
        
        print(f"   有效记录: {len(face_data)} 条")
        
        # 4. 重建索引
        print(f"\n🔨 开始重建索引...")
        faiss_manager.rebuild_from_database(face_data)
        
        # 5. 验证结果
        print(f"\n✅ 索引重建完成！")
        print(f"\n📊 重建后索引状态:")
        print(f"   索引向量数量: {faiss_manager.index.ntotal}")
        print(f"   映射记录数量: {len(faiss_manager.face_id_map)}")
        print(f"   状态: {'✅ 正常' if faiss_manager.validate_index() else '❌ 异常'}")
        
        # 6. 显示映射详情
        print(f"\n📋 映射详情:")
        for i, face_id in enumerate(faiss_manager.face_id_map[:10]):  # 显示前10条
            print(f"   [{i}] -> {face_id}")
        if len(faiss_manager.face_id_map) > 10:
            print(f"   ... 还有 {len(faiss_manager.face_id_map) - 10} 条记录")
        
        print(f"\n🎉 索引重建成功！现在可以正常使用人脸识别服务了。")
        
    except Exception as e:
        app_logger.error(f"重建索引失败: {e}", exc_info=True)
        print(f"\n❌ 重建失败: {e}")
        sys.exit(1)


def check_index():
    """仅检查索引状态，不重建"""
    print("=" * 80)
    print("FAISS索引状态检查")
    print("=" * 80)
    
    print(f"\n📊 索引信息:")
    print(f"   索引文件: {faiss_manager.index_path}")
    print(f"   映射文件: {faiss_manager.face_id_map_path}")
    print(f"   索引维度: {faiss_manager.dimension}")
    
    print(f"\n📈 当前状态:")
    print(f"   索引向量数量: {faiss_manager.index.ntotal if faiss_manager.index else 0}")
    print(f"   映射记录数量: {len(faiss_manager.face_id_map)}")
    
    is_valid = faiss_manager.validate_index()
    if is_valid:
        print(f"   状态: ✅ 正常")
    else:
        print(f"   状态: ❌ 不一致（需要重建）")
        print(f"\n💡 运行以下命令重建索引:")
        print(f"   python scripts/rebuild_faiss_index.py rebuild")
    
    # 从数据库获取统计
    try:
        with get_sync_session() as session:
            face_repo = OrmFaceRepository(session)
            all_faces = face_repo.get_all_faces(include_deleted=False)
            db_count = len(all_faces)
            
            print(f"\n📊 数据库统计:")
            print(f"   人脸记录数量: {db_count}")
            
            if db_count != faiss_manager.index.ntotal:
                print(f"   ⚠️  数据库记录与索引不一致！")
                print(f"   建议重建索引: python scripts/rebuild_faiss_index.py rebuild")
    except Exception as e:
        print(f"\n⚠️  无法连接数据库: {e}")


def main():
    """主函数"""
    if len(sys.argv) > 1:
        command = sys.argv[1].lower()
        if command == "rebuild":
            rebuild_index()
        elif command == "check":
            check_index()
        else:
            print(f"未知命令: {command}")
            print_usage()
    else:
        print_usage()


def print_usage():
    """打印使用说明"""
    print("=" * 80)
    print("FAISS索引管理工具")
    print("=" * 80)
    print("\n使用方法:")
    print("  python scripts/rebuild_faiss_index.py <命令>")
    print("\n可用命令:")
    print("  check   - 检查索引状态")
    print("  rebuild - 重建索引（从数据库）")
    print("\n示例:")
    print("  python scripts/rebuild_faiss_index.py check")
    print("  python scripts/rebuild_faiss_index.py rebuild")
    print("\n💡 提示:")
    print("  - 当出现 'list index out of range' 错误时，运行 rebuild")
    print("  - 重建索引会从数据库读取所有人脸数据")
    print("  - 重建过程中服务可能受影响，建议停止服务后操作")


if __name__ == "__main__":
    main()

